Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy - #24859
Open
adriangb wants to merge 5 commits into
Open
Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy#24859adriangb wants to merge 5 commits into
adriangb wants to merge 5 commits into
Conversation
github-actions
Bot
added
optimizer
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24859 +/- ##
========================================
Coverage 81.72% 81.72%
========================================
Files 1127 1127
Lines 416273 416598 +325
Branches 416273 416598 +325
========================================
+ Hits 340188 340469 +281
- Misses 56094 56110 +16
- Partials 19991 20019 +28 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
run benchmark clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (f47c045) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (f47c045) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
adriangb
added a commit
to pydantic/datafusion
that referenced
this pull request
Sep 1, 2026
The two grouped `COUNT(DISTINCT <string>)` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. apache#24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under apache#24859. Verified from the physical plan with apache#24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged.
Contributor
Author
|
run benchmark clickbench_partitioned clickbench_partitioned clickbench_partitioned env: |
github-actions
Bot
added
logical-expr
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (d510dd4) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
adriangb
added a commit
that referenced
this pull request
Sep 3, 2026
The two grouped `COUNT(DISTINCT <string>)` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. #24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under #24859. Verified from the physical plan with #24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged.
pull Bot
pushed a commit
to TCeason/arrow-datafusion
that referenced
this pull request
Sep 6, 2026
…TINCT) group (apache#24857) A grouped `COUNT(DISTINCT <string>)` over 4,000 groups holding 2 short strings each needs a 36 MB memory budget. It needs 2.0 MB after this change. Every group gets its own hash table, and each table is allocated at warm-up size before the group holds anything, so the memory the query needs tracks the number of groups rather than the amount of data. The query also reports less memory than it holds, so a memory limit does not stop it at the right point. ## Reproduction This needs only `datafusion-cli`. There is no patch, no custom allocator and no data file. ```sql -- repro.sql SET datafusion.execution.target_partitions = 1; -- 4,000 groups with 2 distinct short strings in each. -- avg() stops SingleDistinctToGroupBy from rewriting the distinct aggregate away. SELECT g, count(DISTINCT s) AS d, avg(p) AS a FROM ( SELECT v % 4000 AS g, 'v' || CAST(v AS VARCHAR) AS s, v AS p FROM generate_series(0, 7999) AS t(v) ) GROUP BY g ORDER BY g LIMIT 3; ``` ``` datafusion-cli -m 8M -f repro.sql ``` `s` is a `Utf8View` column, so this exercises `ArrowBytesViewMap`. The 8,000 rows arrive in one batch, so the aggregate builds all 4,000 accumulators before it can emit or spill. Current `main` at `20d1c56761` fails: ``` Resources exhausted: Additional allocation failed for SingleHashAggregateStream[0] with top memory consumers (across reservations) as: DataFusion-Cli#1(can spill: false) consumed 0.0 B, peak 0.0 B, SingleHashAggregateStream[0]#2(can spill: true) consumed 0.0 B, peak 48.0 B, TopK[0]#3(can spill: false) consumed 0.0 B, peak 0.0 B. Error: Failed to allocate additional 111.1 MB for SingleHashAggregateStream[0] with 0.0 B already allocated for this reservation - 8.0 MB remain available for the total memory pool: greedy(used: 0.0 B, pool_size: 8.0 MB) ``` This branch returns the rows: ``` +---+---+--------+ | g | d | a | +---+---+--------+ | 0 | 2 | 2000.0 | | 1 | 2 | 2001.0 | | 2 | 2 | 2002.0 | +---+---+--------+ 3 row(s) fetched. Elapsed 0.006 seconds. ``` Those 4,000 accumulators hold 8,000 short strings, which is about 100 KB of data. The base asks the pool for 111.1 MB to hold it. This branch runs the same query inside `-m 3M`. Both builds return the same rows, and the base does so at `-m 200M`. Each run takes well under a second, and the outcome repeats exactly over three runs on each side. ## How much it improves The minimum memory limit at which that query completes, bisected on each side: | value column | before | after | | --- | --- | --- | | `Utf8` | fails 34 MB, passes 36 MB | fails 1.8 MB, passes 2.0 MB | | `Utf8View` | fails 120 MB, passes 124 MB | fails 2.4 MB, passes 2.6 MB | `clickbench_extended` at `DATAFUSION_RUNTIME_MEMORY_LIMIT: 4G`, pool peak over six runs: | query | base | this branch | change | | --- | --- | --- | --- | | Q2, grouped, 4 string distincts | 98.1 to 98.6 MiB | 11.7 MiB in all six runs | -88.1% | | Q1, ungrouped string distincts | 3.4 MiB | 2.4 MiB | -29.4% | | Q0, ungrouped, high cardinality | 796.8 to 834.8 MiB | 846.5 to 885.5 MiB | +6.3% | Q2 is the only query in any benchmark suite that puts a grouped `COUNT(DISTINCT)` on a non-integer column. Q0 costs more, and it is the one disclosed cost of this PR. Those extra bytes are memory Q0 always held and the pool could not see, not new allocation; appendix A has the three-build decomposition that separates the two. Latency does not move anywhere, which is what an allocation-sizing change should do. ## Which issue does this PR close? No existing issue. I found this when I investigated a production out of memory. I can file an issue if you want it in the changelog. ## Rationale for this change Pre-allocating is right for the one long-lived map behind a `GROUP BY` on a string column. It is wrong for the distinct-count accumulators, because `GroupsAccumulatorAdapter` creates one accumulator per group and most groups hold a handful of values. There the warm-up dwarfs the data. Both maps also under-report the table they hold. `ArrowBytesViewMap` left the control bytes out. `ArrowBytesMap` charged the table only when it grew, so a map that stayed inside its pre-allocation reported its table as free forever. A memory limit acted on a number that was too small. `clear_shrink` is the third part. The aggregate stream calls it to hand memory back before it spills and before a downstream sort. It restored the warm-up capacity instead of releasing it, so nothing came back. ## What changes are included in this PR? - `new` on both maps allocates nothing. A new `with_capacity` keeps the previous pre-allocating behavior. `GroupValuesBytes` and `GroupValuesBytesView` use `with_capacity`, and the two distinct-count accumulators use `new`. A map remembers how it was built, so `take` warms it back up the way it started. - `size()` reports `HashTable::allocation_size()`, the real hashbrown allocation including the control bytes, in place of the old estimate. - A new `clear_and_release` drops every allocation the map holds, and `clear_shrink` calls it. - The value buffer rounds each growth up to a power of two. A lazily grown buffer and a pre-allocated one then sit on one ladder, so a lazy map is never the larger of the two for the same values. Growth stays geometric. - `benches/arrow_bytes_map.rs` moves to `with_capacity` so it keeps measuring the pre-allocating constructor. ## What is the testing strategy for this PR? Two tests in `datafusion/core/tests/memory_limit/mod.rs`, `group_by_count_distinct_utf8` and `group_by_count_distinct_utf8_view`, turn the headline claim into a pass or a fail rather than a number. They run the reproduction query over 4,000 groups with spilling disabled and `target_partitions` pinned to 1, so completing means the query fits the budget rather than spills out of it. The limits are 8 MB and 16 MB, at least 4x clear of both cliffs in the table above. Both tests fail on the merge base and pass here, over five consecutive runs. The `avg(payload)` in the query is load bearing; appendix C says why. Unit tests cover what the memory-limit tests cannot see: that `new` allocates nothing, that `with_capacity` reports a table size bracketed by an independently derived lower bound, that `take` preserves the configured capacity, that `clear_shrink` drops the reported size to near zero, and that a lazily grown buffer never exceeds a pre-allocated one holding the same values. Run locally on the rebased head, all passing: `datafusion-physical-expr-common` (87 lib, 8 doc), `datafusion-functions-aggregate-common` (49), `datafusion-functions-aggregate -- count_distinct` (2), `datafusion-physical-plan -- group_values` (96) and the `memory_limit` module (39, which includes the `count_distinct_spill` test that arrived on `main` in apache#24888 and apache#24918). `cargo fmt --check` and `cargo clippy --all-targets -D warnings` are clean on the changed crates. CI has not yet run this branch against the new base. No query results change. ## Are there any user-facing changes? Yes, in `datafusion-physical-expr-common`. `ArrowBytesMap::new` and `ArrowBytesViewMap::new` no longer pre-allocate, and callers that want the previous behavior should use the new `with_capacity`. Both types also gain `clear_and_release`. This changes an existing public constructor rather than adding one, so tell me if you would like the `api change` label. For users, a grouped `COUNT(DISTINCT)` on string and binary columns uses much less memory and reports its usage to the `MemoryPool` accurately. A query that previously hit a memory limit may now succeed. --- ## Appendix A: Query 0 costs 6.3% more Q0 is `COUNT(DISTINCT)` over three high-cardinality strings with no `GROUP BY`. It is a handful of maps that each grow to millions of entries, which is the opposite population from the one this PR targets. The pre-allocation was never the dominant cost there, so removing it buys nothing. Over six runs the base spans 796.8 to 834.8 MiB and this branch spans 846.5 to 885.5 MiB. The ranges do not overlap, so the effect is real and not run-to-run noise. Three local builds on a deterministic subset separate the two changes. The middle build differs from the base only in the accounting, because restoring the warm-up makes the constructors byte-identical to base: | build | Q0 pool peak | | --- | --- | | base | 48,421,820 | | this branch with the warm-up restored | 51,048,396 | | this branch | 51,018,828 | That decomposes the increase exactly. +2,626,576 is the accounting correction: `allocation_size()` charges the real hashbrown allocation, which is `4 * buckets + 5,384` more than the old formula, being the control bytes plus the 7/8 load-factor slack. -29,568 is the lazy constructor, which makes Q0 slightly better. Reverting the accounting would restore an under-report of about 19% on this path. That under-report is the bug this PR exists to fix, and the memory-limit result above depends on fixing it. ## Appendix B: what one accumulator costs One per-group accumulator holding a single 24-byte value: | | before, actual | before, reported | after | | --- | --- | --- | --- | | `BytesDistinctCountAccumulator` | 14,648 B | 8,240 B | 180 B | | `BytesViewDistinctCountAccumulator` | 33,920 B | 28,792 B | 260 B | The middle column is the reporting gap. The `Utf8` map really held 14,648 bytes and reported 8,240, because the whole hash table was invisible to the old accounting. These are measured directly rather than asserted in a test, since the exact numbers follow the hashbrown layout. ## Appendix C: notes on the tests and the benchmarks The memory-limit query uses `avg(payload)`, not `count(*)`. A non-distinct `count` lets `SingleDistinctToGroupBy` rewrite the distinct aggregate into a plain two-stage `GROUP BY`. The per-group accumulators would then never exist, and the tests would pass by construction on the base commit too. That rule accepts a non-distinct `sum`, `min` or `max` because each re-aggregates its own partial results correctly over the deduplicated inner group by. `avg` does not, so the rule can never admit it under any extension, including the one apache#24859 proposes. The benchmark figures were measured against the previous merge base `da89c7c85b`. They have not been re-run against the current base `20d1c56761`. The commits after `84f07da` on this branch touch only `datafusion/core/tests/memory_limit/mod.rs`, so nothing on this branch since then can move a benchmark, but the base itself has moved. Pool peak is the instrument here, not peak RSS. Pool peak reproduces to under 1% on a 98 MiB query and exactly on the 3.4 and 11.7 MiB ones. Peak RSS on this harness has a 4.1% standard deviation over 11 readings of identical code plus a 2.3% order bias, and shows no effect from this change once that null is accounted for. ## Follow-ups, not in this PR - The same undercount remains at five other production `insert_accounted` call sites: `group_values/row.rs:171`, `multi_group_by/mod.rs:434,554`, `multi_group_by/dictionary.rs:197,584` and `array_agg.rs:989`. Each is one map per query, so the absolute error is bounded, and the fix is the same one-line swap. - The `count_distinct_groups` benchmarks in `datafusion/functions-aggregate/benches/count_distinct.rs` cover `Int64`, `Int32` and `UInt32` only, so this path has no criterion coverage. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
adriangb
force-pushed
the
claude/single-distinct-to-groupby-allow-count
branch
from
September 7, 2026 06:21
d3dbc09 to
e0e5a8c
Compare
adriangb
marked this pull request as ready for review
September 7, 2026 07:10
Contributor
Author
|
run benchmark clickbench_extended clickbench_extended clickbench_extended |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedCPU Details (lscpu)Details
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedCPU Details (lscpu)Details
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (3d3cf6b) to e1ca94f (merge-base) diff Run configurationrun benchmark clickbench_extendedCPU Details (lscpu)Details
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
adriangb
force-pushed
the
claude/single-distinct-to-groupby-allow-count
branch
from
September 7, 2026 14:18
3d3cf6b to
e0e5a8c
Compare
Contributor
Author
`SingleDistinctToGroupBy` rewrites `AGG(DISTINCT x)` into a two phase group by, which is what keeps a high cardinality distinct off the one-accumulator-per-group path in `GroupsAccumulatorAdapter`. The rule tolerated a non-distinct `sum`, `min` or `max` next to the distinct aggregate but bailed out on `count`, so the very common `count(*), count(DISTINCT x) ... GROUP BY` shape kept the unrewritten plan and its memory profile. Allow a non-distinct `count` as well. `count` is the one supported function whose outer phase is a different function: the inner group by counts the rows of each `(group, distinct value)` partition and the outer phase adds those partial counts up with `sum`, since count over a group is the sum of the counts of any partition of that group. Two details follow from that substitution: - `count` and `sum` are resolved from the session function registry and the rewrite only fires when the aggregate is that exact `count`, so a session without a registry or with its own `count` is left alone. - `count` returns a non-null 0 over an empty input while `sum` of no rows is NULL, which is reachable for an aggregate with no group by. The projection selects `CASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END`, which restores the 0 and keeps the column's type and nullability as `count` had them. The new sqllogictest file asserts every result twice, once with the optimizer disabled and once with it enabled, over data with NULL and all-NULL distinct values, an empty input, and `count(*)` versus `count(col)` versus `count(1)`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`aggregate_distinct_with_having` round trips `SELECT a, count(distinct b) ... HAVING count(b) > 100` through substrait and asserts the plan comes back displaying identically. It passed only because the non-distinct `count` made `SingleDistinctToGroupBy` bail out, so the plan had no aliases in it. With the rule now allowing that `count`, the query is rewritten and the assertion fails. The failure is a pre-existing substrait gap rather than anything specific to this query: substrait carries no names for an aggregate's grouping and measure expressions, so the consumer derives them from the expressions themselves and the `alias1` and `alias2` names the rule introduces are lost. Any plan the rule rewrites fails the same way, including the plain `SELECT a, count(distinct b) FROM data GROUP BY a, c` that this change does not touch. Remove the rule from the session used by this one test, so it keeps covering the un-rewritten aggregate it was written for instead of depending on the rule bailing out. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A memory limited ClickBench run contradicted the memory rationale for the
case it measures. Q22 is `SELECT "SearchPhrase", MIN("URL"), MIN("Title"),
COUNT(*), COUNT(DISTINCT "UserID") ... GROUP BY "SearchPhrase"`, the only
plan the change touched, and its peak memory pool reservation went from
3.1 MiB to 7.3 MiB, up 132.2%, with neighbouring queries moving by a few
percent either way.
`UserID` is an `Int64`, and `Count::groups_accumulator_supported` returns
true for every integer type and false for everything else. So the base
plan already had `PrimitiveDistinctCountGroupsAccumulator` and never went
near `GroupsAccumulatorAdapter`. The rewrite replaced a compact vectorized
accumulator with an inner group by on `(SearchPhrase, UserID)`, which also
carries `min("URL")` and `min("Title")` at that much finer grain, and
bought nothing back.
Measured locally on 4M rows, 500k groups and 2M distinct pairs, peak pool
reservation for `SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g`:
x BIGINT 135 MiB unrewritten, 205 MiB rewritten
x VARCHAR 13.6 GiB unrewritten, 255 MiB rewritten
Peak RSS for the VARCHAR pair was 18.0 GiB against 433 MiB, so the pool
figure is real memory rather than an accounting artifact. The rewrite pays
exactly when the distinct argument would otherwise land on the adapter.
`datafusion-optimizer` cannot depend on `datafusion-functions-aggregate` to
read that list of types, and has no `AccumulatorArgs` to ask
`groups_accumulator_supported` with. Its `Cargo.toml` names the way out, so
this adds `AggregateUDFImpl::groups_accumulator_supported_for_types`,
defaulting to false as `groups_accumulator_supported` does. `Count` is the
only implementor and its physical method now delegates to it, so the two
cannot disagree.
The gate covers only the `count` this branch added. A plan that already
qualified through a non-distinct `sum`, `min` or `max` is rewritten exactly
as before, over any distinct argument type.
The ClickBench snapshot returns to its base form, so no existing snapshot
in the repository changes. `single_distinct_to_groupby.slt` now carries the
same values in a `VARCHAR` and an `INT` column and asserts both sides of
the gate, still under both `datafusion.optimizer.max_passes = 0` and the
default.
`cargo doc` runs with `-D warnings`, and a doc link from the public `SingleDistinctToGroupBy` to the private `rewrite_pays_for_count` is an error. Say the same thing in prose instead.
`AggregateUDFImpl::groups_accumulator_supported_for_types` returned `bool` with a `false` default, and `rewrite_pays_for_count` read `false` as "the rewrite pays". Every aggregate that did not override the method was therefore waved through the gate, which is the permissive answer rather than the safe one. Return `Option<bool>` instead, defaulting to `None`, and rewrite only on `Some(false)`. `None` says the aggregate does not answer the question, and silence is not evidence that the rewrite pays. Measured on unmodified upstream, over 4,000,000 rows in 2,000 groups, `SELECT g, count(*), sum(DISTINCT int_col) FROM t GROUP BY g` reached the gated path and regressed 3.15x: 71.0 MiB unrewritten against 223.4 MiB rewritten. `min(DISTINCT x)` reached it too, and regresses much further, because `min(DISTINCT x)` is `min(x)` and the unrewritten plan keeps one scalar per group. `Count` is the only implementor and its answers are unchanged, so no plan that the gate already allowed changes shape.
adriangb
force-pushed
the
claude/single-distinct-to-groupby-allow-count
branch
from
September 7, 2026 17:10
e0e5a8c to
07f24b6
Compare
adriangb
added a commit
to pydantic/datafusion
that referenced
this pull request
Sep 7, 2026
…over a string (apache#25026) ## Which issue does this PR close? No existing issue. This is benchmark coverage carved out of apache#24859 so that it can land first: a benchmark query added in the same PR that needs it cannot appear in an A/B run, because the bot compares against the merge base and the merge base does not have the query. ## Rationale for this change `COUNT(DISTINCT)` has a specialized `GroupsAccumulator` for the integer types and for no other type. Every other type falls back to `GroupsAccumulatorAdapter`, which holds one boxed `Accumulator`, and therefore one hash table, for each group. The cost of that fallback is per group, so the group cardinality is what decides how much it costs. Nothing in either suite measures that: - Extended Q2 is the only query that puts a `COUNT(DISTINCT)` on a non-integer column at all. It groups by `BrowserCountry`, so it exercises the adapter at a low group cardinality, which is where the adapter is cheapest. - Standard Q8, Q10, Q11 and Q13 hold a distinct aggregate that stands alone, so `SingleDistinctToGroupBy` already rewrites them and they never reach the adapter. - Standard Q9 carries an `AVG`, which that rule has never accepted. - Standard Q22 is the one query that pairs a lone distinct aggregate with a non-distinct count, and it counts distinct `UserID`, an `Int64`, which has a specialized accumulator. So a lone `COUNT(DISTINCT <string>)` grouped by a high cardinality key, next to a non-distinct `COUNT(*)`, is uncovered. That is an ordinary analytics shape, it is the shape apache#24857 changed the memory profile of, and it is the shape apache#24859 proposes to rewrite. Neither of those could show its effect on any benchmark in this repository. ## What changes are included in this PR? The query: ```sql SELECT "SearchPhrase", COUNT(*) AS c, COUNT(DISTINCT "MobilePhoneModel") AS models FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; ``` The extended ClickBench queries live in three places, and a query added to only one of them is invisible to the others. This PR adds it to all three: - `benchmarks/queries/clickbench/extended/q14.sql`, which feeds `dfbench clickbench --queries-path`. Queries are discovered from the directory, so the runner needs no change. - `benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q14.benchmark`, which feeds `benchmark_runner clickbench_extended`. Structurally identical to `q13.benchmark` apart from the query. Suite files are also discovered from the directory; `benchmark_runner clickbench_extended --list` now reports 15 queries. - `datafusion/sqllogictest/test_files/clickbench_extended.slt`, which runs the extended queries against the committed ten row `clickbench_hits_10.parquet` fixture. That last file had only ever mirrored q0 through q6. q7 through q13 were added to the queries directory without a matching sqllogictest entry, and q14 would have widened that gap, so this backfills q7 through q14 together. That is why the diff is larger than one query. Having to add one query in three places is itself the problem, and the copies have already drifted: extended q6 carries a cast in its `sql_benchmarks` copy that the other two do not have. I filed apache#25031 for that; it is out of scope here. There is one pre-existing staleness this PR does not fix. `datafusion/core/benches/sql_planner.rs` builds its ClickBench planning set from a hardcoded `(0..=7)` over the extended directory, so extended Q8 through Q13 were already outside it before this PR and Q14 joins them. That is a separate cleanup. ## What is the testing strategy for this PR? The sqllogictest entries are the test. `clickbench_extended.slt` runs every extended query against the committed ten row fixture and asserts its output, so q7 through q14 are now executed on every CI run rather than only by whoever runs the benchmark suite by hand. `cargo test -p datafusion-sqllogictest --test sqllogictests -- clickbench` passes, 2 of 2 files. `cargo test -p datafusion-benchmarks` passes, 192 tests, and `benchmark_runner clickbench_extended --list` reports the suite at 15 queries, confirming the new `.benchmark` file parses and is discovered. I do not have a `hits.parquet` to hand, so I have not run the query against the full dataset. The plan shape, which is the whole value of the query, was checked separately: on an equivalent local table with the same filter, grouping, `COUNT(*)` and string `COUNT(DISTINCT)`, the aggregate keeps `aggr=[[count(Int64(1)), count(DISTINCT ...)]]`, which is the `GroupsAccumulatorAdapter` path this query exists to measure, and it is not rewritten away. What I have not established is the effect size on the real dataset, since that depends on how many distinct `SearchPhrase` values survive the filter. If a reviewer with the data runs it, that number is worth having on this PR. `ci/scripts/doc_prettier_check.sh` passes on the README change. ## Are there any user-facing changes? No. This adds a benchmark query and documentation only. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Contributor
Author
|
run benchmark clickbench_extended clickbench_extended clickbench_extended env: |
Contributor
Author
|
show benchmark queue |
|
Hi @adriangb, you asked to view the benchmark queue (#24859 (comment)).
File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (07f24b6) to 16ace4f (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "16G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
No existing issue. We found this during an investigation of a production out of memory. We can file one if you want a changelog entry.
Rationale for this change
SingleDistinctToGroupByrewritesAGG(DISTINCT x)into a two phase group by, which is what keeps a high cardinality distinct off the one-accumulator-per-group path inGroupsAccumulatorAdapter. The rule accepts a non-distinctsum,minormaxnext to the distinct aggregate. It rejects a non-distinctcount. Onecount(*)is therefore enough to keep the unrewritten plan, andcount(*), count(DISTINCT x) ... GROUP BYis a very common shape.This PR allows that
count, gated on the distinct argument having no specializedGroupsAccumulator.Reproduce it
Run this in
datafusion-cli. It writes 4,000,000 rows in 500,000 groups, with 2,000,000 distinct(g, x)pairs.On
mainthe rule leaves the aggregate alone:This branch rewrites it:
The results are identical.
Why the rewrite helps here
The rewrite takes the distinct aggregate off
GroupsAccumulatorAdapter, which keeps one boxedAccumulator, and therefore one hash table, for every group.count(DISTINCT x)has a specializedGroupsAccumulatorfor each integer type and for no other type, so a string argument is exactly the case that lands on the adapter.The rewrite is not free. Every other aggregate moves down into the inner group by, which holds one row for each
(group, distinct value)pair rather than one row for each group, and each aggregate then keeps its state at that finer grain.There is a second effect that only appears with more than one partition. The unrewritten plan's partial aggregate sees every group in every partition, so its per-group accumulators are duplicated
target_partitionstimes. The rewritten plan hash-partitions the inner group by, so the distinct values are split across partitions instead of replicated. This is why the gap widens rather than narrows as partitions go up.The measurements
An earlier version of this description carried two columns, one for
mainand one for a tree that also carried #24857. #24857 is merged, so there is one baseline again and every number below is new.Base is
e1ca94fb11, which contains #24857 asd25ffaab49. The branch is this PR rebased onto it. Both sides are release builds of the same harness: it wraps aGreedyMemoryPoolof 64 GiB, which is never reached, in the in-treePeakRecordingPool, registers one parquet file, streams the result without collecting it, and reports the pool high-water mark andgetrusagepeak RSS. Every figure is the median of three runs. Attarget_partitions = 1the three runs agreed to the byte in every cell; at 8 they agreed to within 1%, except two branch cells at 7% and 19%.Every shape is 4,000,000 rows and the query is
SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g.Each file carries the distinct argument three times, as
Utf8, asUtf8Viewand asBIGINT, written with an explicitarrow_castso the type is the file's and not a reader setting. Thexin the reproduction above isUtf8View, sincedatafusion.execution.parquet.schema_force_view_typesdefaults to true; it is the 500,000 groupUtf8Viewrow of the first table.target_partitions = 8Utf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewBIGINTtarget_partitions = 1Utf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewUtf8Utf8ViewThree things in there are worth calling out.
The wall clock was not the point of this PR and it is the largest effect in the table. At
target_partitions = 8and 500,000 groups the query goes from about a second to about 40 ms. That is the accumulator-per-group construction cost, and it grows with partitions on the unrewritten side while the rewritten side parallelizes.Peak RSS improves in every measured shape, including the two where the pool peak regresses. The unrewritten plan's RSS sits far above what it reserves; millions of small independent hash tables fragment in a way one large table does not.
The
BIGINTrow is the control for the gate. The plan is identical on both sides, and so is every measurement.Where the rewrite stops paying
The crossover is in the density of distinct values, not in the group count. It arrives only when nearly every row holds a distinct value, and the loss there is bounded: 1.30x more peak pool for
Utf8and 1.10x forUtf8View, at 2,000 groups over 4,000,000 fully distinct values with 8 partitions, and a wash at one partition. Those same two cells use about half the RSS and run four times faster.At 500,000 groups over 4,000,000 fully distinct values, which is the same density, the rewrite is 1.61x to 1.96x better again, because the unrewritten side now also pays for 500,000 accumulators.
The gate is a proxy
The gate asks which accumulator the distinct aggregate gets. That is not the true discriminator.
What decides the outcome is the cost per distinct value on each side. The rewrite materializes one hash table row for each distinct
(group keys, x)pair plus one accumulator slot for each companion aggregate at that grain. It wins when the unrewritten accumulator costs more than that per value, and loses when it costs less.The proxy and the true discriminator agree for
count(DISTINCT x), which is the only case this PR opens. They disagree elsewhere, and the pre-existing arms of this rule carry that disagreement:Re-measured on the same base as the tables above, over 4,000,000 rows in 2,000 groups with 2,000,000 distinct values, comparing the rewritten plan against the same query with a
count(*)added to hold the rule off:sum(DISTINCT int_col)andavg(DISTINCT int_col)have no distinct groups accumulator, andmainrewrites both today. Both regress: 2.67x more peak pool at one partition, 1.79x at eight.min(DISTINCT x)is far worse.min(DISTINCT x)is the same value asmin(x), andmin_maxcorrectly ignoresis_distinct, so the unrewritten plan holds one scalar for each group while the rewrite builds a hash table over every distinct pair. That is 543x more peak pool at one partition and 126x at eight.That regression predates this PR, and this PR does not extend it. The gate keeps every one of those functions out of the path this PR adds. We report the broader regression upstream separately. A cost model is out of scope here.
What changes are included in this PR?
This PR allows a non-distinct
countnext to the distinct aggregate, when the distinct aggregate reports that it has no specializedGroupsAccumulatorfor its argument types.countis the one supported function whose outer phase is a different function. The inner group by counts the rows of each(group, distinct value)partition. The outer phase adds those partial counts withsum, because acountover a group is the sum of the counts of any partition of that group.Two details follow from that substitution.
countandsumcome from the session function registry, asreplace_distinct_aggregatealready does forfirst_value. The rewrite fires only for that exactcount, compared by identity and not by name. A session without a registry, or with its owncount, keeps the previous behaviour.countreturns a non-null0over an empty input, andsumof no rows returns NULL. An aggregate without aGROUP BYreaches that case. The inner aggregate emits no rows and the outer aggregate still emits one row, soSELECT count(*), count(DISTINCT x) FROM emptywould returnNULL, 0instead of0, 0. The projection selectsCASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END. That restores the0, and it keeps the type and the nullability thatcounthad.FILTERandORDER BYstill block the rewrite.The gate
An optimizer rule has no
AccumulatorArgsto askAggregateUDFImpl::groups_accumulator_supportedwith.datafusion-optimizeralso cannot depend ondatafusion-functions-aggregateto readcount's list of types.datafusion/optimizer/Cargo.tomlnames the intended way out:So this PR adds
AggregateUDFImpl::groups_accumulator_supported_for_types(&[DataType], is_distinct) -> Option<bool>.Countis the only implementor.Count::groups_accumulator_supportednow delegates to it, so there is one list of supported types and not two that can drift.The default is
None, which means the implementation does not answer the question.Noneis not a third answer, and a caller must not read it as eitherSome(true)orSome(false). This rule rewrites only onSome(false), which is the one answer that reports a call on the adapter. An aggregate that answersNone, which today is every aggregate exceptcount, keeps the previous behaviour.The gate covers only the
countthis PR adds. A plan that already qualifies through a non-distinctsum,minormaxis rewritten as before, over any distinct argument type. A narrower rule there would change plans that this repository has always rewritten, and nothing measured here calls for that change.Files touched beyond the rule itself:
datafusion/expr/src/udaf.rsanddatafusion/functions-aggregate/src/count.rs: the new trait method and its one implementation.datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt: the new coverage described below.datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:aggregate_distinct_with_havingnow builds its session without this rule, so it keeps round tripping the plan shape that the test was written for. Its distinct argument is aDECIMAL, so the gate does not exclude it.No existing snapshot in the repository changes.
What is the testing strategy for this PR?
single_distinct_to_groupby.sltasserts every result twice, once underdatafusion.optimizer.max_passes = 0and once under the default, with identical expected blocks. A null-handling error or a type error therefore appears as a result mismatch, and not only as a plan difference.The table carries the same values in a
VARCHARcolumn and in anINTcolumn, which are the two sides of the gate. The file asserts both sides. TheVARCHARdistinct rewrites with acountnext to it. TheINTdistinct does not. TheINTdistinct still rewrites when it qualifies throughsum. The results agree with the unoptimized plan in every case.The file also asserts that an aggregate which answers
Nonestays unrewritten.sum(DISTINCT v)andmin(DISTINCT v)next to acount(*)both keep the plan they have onmain.The remaining coverage is
count(*)againstcount(1)againstcount(col), grouped and ungrouped, a group whose distinct column is entirely NULL, a group with NULLs in both the distinct column and the summed column, empty input in three shapes,HAVINGwithORDER BYon the rewritten count, and the production join shape.The unit tests of the rule cover both sides of the gate directly. They also cover the
Noneanswer twice: once withsum(DISTINCT b), and once with a test aggregate that leaves the new method at its default.Run locally on the rebased head, all passing:
datafusion-optimizer,datafusion-expranddatafusion-functions-aggregatelib and integration tests, theaggregate_distinct_with_havingsubstrait roundtrip, and the whole sqllogictest suite at 511 of 511 files.Outside the tests, the harness above ran
SELECT sum(a), sum(b), count(*) FROM (SELECT g, count(*) AS a, count(DISTINCT x) AS b FROM t GROUP BY g)on both sides for three of the shapes in the table. The two sides agree in every case.Benchmarks
clickbench_extendedQ14 landed onmainin #25026 and is the one query in any suite with this shape: a loneCOUNT(DISTINCT <string>)next to a non-distinctCOUNT(*), grouped by a high cardinality string. Three runs of this branch against merge base16ace4f,DATAFUSION_RUNTIME_MEMORY_LIMIT: 16G(trigger):The memory figure reproduces to a tenth of a percent across the three. This is 100,000,000 rows of real ClickBench data, and it is a larger effect than the synthetic measurements above, not a smaller one.
The limit is 16G rather than the 4G used on earlier runs of this PR because the base side of Q14 peaks at 4.4 GiB. At 4G the baseline runs out of memory and there is no comparison to report.
Run 3 was on a contended machine: its suite total is 42.1 s against 30.8 s and 31.0 s for the other two, and it reports Q0 at 1.30x slower, Q7 at 1.14x and Q8 at 1.13x. None of those three can change plan here. Extended Q0 has three distinct arguments, so
fields_set.len() != 1. Extended Q7 and Q8 have no distinct aggregate at all. They are an in-experiment control, and they say what the noise floor of that run was.No other query changes plan, and the changed-plan set is empty by construction rather than by measurement. The rule needs exactly one distinct argument and accepts only
sum,min,maxand nowcountas companions:extendedq0, q1, q2fields_set.len() != 1mainalready rewrites themAVG, which the rule has never acceptedCOUNT(DISTINCT "UserID")is over anInt64, which has a specializedGroupsAccumulatorThat holds in the measurements. Outside Q14, every query in all three runs reports
no changeon time, and the memory table is+0.0%or-0.1%except Q10 and Q12, which move by fractions of a MiB on 2.5 MiB and 1.4 MiB reservations. TheBIGINTrow of the table above is theq22case measured directly, at 1.00x on both the pool peak and the wall clock.An earlier version of this description reported a
+132.2%memory regression onq22. That was the ungated version of this branch. The gate removed it, andq22's plan is now byte for byte the base plan.Are there any user-facing changes?
There is no public API change and no change to query results.
AggregateUDFImplgains one method with a default, which is not a breaking change for implementors.Plans for
SELECT ..., count(...), count(DISTINCT x) ... GROUP BY ...change shape whenxhas no specializedGroupsAccumulator.EXPLAINoutput for that shape therefore differs, and such queries should use less memory and run faster.